use Wx 0.15 qw[:allclasses];
use strict;
package wxSolutionComparison;

use Wx qw[:everything];
use base qw(Wx::Panel);
use strict;
use Storable;
use SetSolution;
use wxTreeWindow;
use wxSortableGrid;
use Scalar::Util qw(blessed);

use constant INDICES => { 'PE' => 'Pendant Edge' , 'ES' => 'Equal Splits', 'FP' => 'Fair Proportion', 'SHAPLEY' => 'Shapley Value'};
use constant SUMMARY => { 'Max' => 'Maximum Value', 'Mean' => 'Mean Value', 'Min' => 'Minimum Value'};

sub new {
	my( $self, $parent, $id, $pos, $size, $style, $name ) = @_;
	    
	$parent = undef              unless defined $parent;
	$id     = -1                 unless defined $id;
	$pos    = wxDefaultPosition  unless defined $pos;
	$size   = wxDefaultSize      unless defined $size;
	$name   = ""                 unless defined $name;

	$style = wxDEFAULT_FRAME_STYLE 
		unless defined $style;

	#Create controls
	$self = $self->SUPER::new( $parent, $id, $pos, $size, $style, $name );
	
	$self->{splitter_window} = Wx::SplitterWindow->new($self, -1, wxDefaultPosition, wxDefaultSize, wxSP_3D|wxSP_BORDER);
	$self->{LeftPanel} = Wx::Panel->new($self->{splitter_window}, -1, wxDefaultPosition, wxDefaultSize);
	$self->{RightPanel} = Wx::Panel->new($self->{splitter_window}, -1, wxDefaultPosition, wxDefaultSize);
	
	#The grid
	$self->{grid_solutions} = wxSortableGrid->new($self->{LeftPanel}, -1, wxDefaultPosition, wxDefaultSize,wxRAISED_BORDER);
	#The tree
	$self->{wxtree} = wxTreeWindow->new($self->{RightPanel}, -1);

	#Objects for adding/removing columns
	$self->{combo_add_index} = Wx::ComboBox->new($self->{LeftPanel}, -1, '', wxDefaultPosition, wxDefaultSize, [],wxCB_READONLY);
	$self->{button_add_index} = Wx::Button->new($self->{LeftPanel}, -1, "Add Index");
	$self->{button_del} = Wx::Button->new($self->{LeftPanel}, -1, "Remove Column");

	#Objects for budget selection
	$self->{label_budget} = Wx::StaticText->new($self->{LeftPanel}, -1, "Budget:", wxDefaultPosition, wxDefaultSize, wxALIGN_RIGHT);
	$self->{text_budget} = Wx::TextCtrl->new($self->{LeftPanel}, -1, "0", wxDefaultPosition, wxDefaultSize);
	$self->{check_all_budget_solutions} = Wx::CheckBox->new($self->{LeftPanel}, -1, "Show all solutions with budget", wxDefaultPosition, wxDefaultSize);

	#Do the layout
	$self->{sizer} = Wx::BoxSizer->new(wxVERTICAL);

	$self->{sizer_option_bar} = Wx::BoxSizer->new(wxHORIZONTAL);
	$self->{sizer_addremove} = Wx::GridSizer->new(2,2,2,2);
	$self->{sizer_budget} = Wx::BoxSizer->new(wxVERTICAL);
	$self->{sizer_budget_entry} = Wx::BoxSizer->new(wxHORIZONTAL);
	
	$self->{sizer}->Add($self->{splitter_window}, 1, wxEXPAND, 0);

	#Add/remove column objects
	$self->{sizer_addremove}->Add($self->{combo_add_index},0,wxEXPAND|wxALL|wxALIGN_CENTER_VERTICAL,0);
	$self->{sizer_addremove}->Add(Wx::Panel->new($self->{LeftPanel},-1),0,wxEXPAND|wxALL|wxALIGN_CENTER_VERTICAL,0);
	$self->{sizer_addremove}->Add($self->{button_add_index},0,wxEXPAND|wxALL|wxALIGN_CENTER_VERTICAL,0);
	$self->{sizer_addremove}->Add($self->{button_del},0,wxEXPAND|wxALL|wxALIGN_CENTER_VERTICAL,0);
	
	#Budget objects
	$self->{sizer_budget_entry}->Add($self->{label_budget});
	$self->{sizer_budget_entry}->Add($self->{text_budget});
	$self->{sizer_budget}->Add($self->{sizer_budget_entry});
	$self->{sizer_budget}->Add($self->{check_all_budget_solutions});
	
	#Option bar
	$self->{sizer_option_bar}->Add($self->{sizer_addremove},0,wxEXPAND|wxALIGN_LEFT);
	$self->{sizer_option_bar}->Add($self->{sizer_budget},0,wxEXPAND|wxALIGN_RIGHT|wxLEFT,40);

	$self->{sizer_top} = Wx::BoxSizer->new(wxVERTICAL);
	$self->{sizer_bottom} = Wx::BoxSizer->new(wxVERTICAL);
	$self->{sizer_top}->Add($self->{grid_solutions}, 1, wxEXPAND, 0);
	$self->{sizer_top}->Add($self->{sizer_option_bar}, 0, wxEXPAND|wxALL, 5);
	$self->{sizer_bottom}->Add($self->{wxtree}, 30, wxEXPAND, 0);
	$self->{splitter_window}->SplitHorizontally($self->{LeftPanel}, $self->{RightPanel}, 200);
	
	$self->{splitter_window}->SetMinimumPaneSize(200);
	
	$self->SetSizer($self->{sizer});
	$self->{LeftPanel}->SetSizer($self->{sizer_top});
	$self->{RightPanel}->SetSizer($self->{sizer_bottom});
	#not really necessary -> but leaving for future reference
	$self->{sizer}->Fit($self);
	
	$self->Layout();
	
	#Register events
	$self->{grid_solutions}->set_selection_handler($self);
	Wx::Event::EVT_COMBOBOX($self, $self->{combo_add_index}->GetId, \&combo_selected);
	Wx::Event::EVT_BUTTON($self, $self->{button_add_index}->GetId, \&add_index);
	Wx::Event::EVT_BUTTON($self, $self->{button_del}->GetId, \&remove_col);
	Wx::Event::EVT_TEXT($self, $self->{text_budget}->GetId, \&budget_change);
	Wx::Event::EVT_CHECKBOX($self, $self->{check_all_budget_solutions}->GetId, \&show_all_budget_solutions);
	$self->{grid_solutions}->set_row_omit(2);
	

	return $self;

}

#Set the state
sub set_state {
	my $self = shift;
	$self->{state} = shift;
}

#Called whenever the state changes
sub update {
	my $self = shift;
	
	$self->update_index_hash();
	$self->set_combo_add_index();
	
	#If the tree or solutions have changed we reset the grid completely and make sure that the tree will be redrawn
	unless (defined $self->{last_tree} && $self->{last_solutions} && 
      $self->{last_tree}==$self->{state}->{tree} && 
	  $self->{last_solutions} == $self->{state}->{solutions}) {
		$self->create_empty_grid();
		delete $self->{last_drawn_index};
	}
	
	#Add any solutions not in the grid to it
	$self->update_grid();

	#Store the last tree and solutions
	$self->{last_tree} = $self->{state}->{tree};
	$self->{last_solutions} = $self->{state}->{solutions};
	
}

#Updates the hash that controls which indices are displayed 
sub update_index_hash {
	my $self = shift;
	
	#Add the species specific indices
	foreach (keys %{+INDICES}) {
		$self->{index}->{$_}->{long} = INDICES->{$_};
		$self->{index}->{$_}->{short} = $_;
		unless ($self->{index}->{$_}->{display}) {
			$self->{index}->{$_}->{display} = 0;
		}
		$self->{index}->{$_}->{type} = 'species_specific';
		$self->{index_r}->{+INDICES->{$_}} = $self->{index}->{$_};
	}

	#Add the summary indices
	foreach (keys %{+SUMMARY}) {
		$self->{index}->{$_}->{long} = SUMMARY->{$_};
		$self->{index}->{$_}->{short} = $_;
		unless ($self->{index}->{$_}->{display}) {
			$self->{index}->{$_}->{display} = 0;
		}
		$self->{index}->{$_}->{type} = 'summary';
		$self->{index_r}->{+SUMMARY->{$_}} = $self->{index}->{$_};
	}

	#Add the named solutions and solutions of appropriate budget
	if (defined $self->{state}->{solutions}) {
		my @sol_named = $self->{state}->{solutions}->get_solution_names();
		my @sol_cost = $self->{state}->{solutions}->get_solution_names_cost($self->{text_budget}->GetValue());
		my @sol = (@sol_named, @sol_cost);
		my %todel;
		#Todel contains all the keys already in $self->{index}
		foreach (keys %{$self->{index}}) { $todel{$_} = 1;}

		#We add each of the solutions
		foreach (@sol) {
			$self->{index}->{$_}->{long} = $_;
			$self->{index}->{$_}->{short} = $_;
			unless ($self->{index}->{$_}->{display}) {
				$self->{index}->{$_}->{display} = 1;
			}
			$self->{index}->{$_}->{named} = 0;
			$self->{index}->{$_}->{type} = 'solution';
			$self->{index_r}->{$_} = $self->{index}->{$_};
			#We remove this key from %todel
			delete $todel{$_}
		}
		#We flag the named solutions as being named
		foreach (@sol_named) {$self->{index}->{$_}->{named} = 1; }
		
		#We delete any keys which don't have corresponding solutions
		foreach (keys %todel) {	if ($self->{index}->{$_}->{type} eq 'solution') {delete $self->{index}->{$_}; }}
	}
}

#Called when the budget changes
sub budget_change {
	my $self = shift;
	#First make sure that solutions of this budget are in the hash
	$self->update_index_hash();
	#Then set the display value for these solutions
	$self->show_all_budget_solutions();
	#Now scale trees (makes appropriate changes to the grid)
	$self->calculate_scaled_trees();
	#Finally display the selected tree
	$self->display_selected();
}

#Called when the "show all solutions with that budget" checkbox is clicked
sub show_all_budget_solutions {
	my $self = shift;

	#We set the solutions display values appropriately
	foreach (values %{$self->{index}}) {
		if ($_->{type} eq 'solution' && $_->{named} == 0) {
			$_->{display} = $self->{check_all_budget_solutions}->GetValue();
		}
	}
	
	$self->update_grid();
	
}

#Calculate scaled trees for indices that are displayed. Also displays the epd and cost for these
#solutions in the grid.
sub calculate_scaled_trees {
	my $self = shift;
	
	#Create a hash for grid rows and columns (for entering epd and cost details)
	my %col_hash;
	my $nrows = $self->{grid_solutions}->GetNumberRows();
	
	for (my $index = 0; $index < $self->{grid_solutions}->GetNumberCols(); $index++) {
		$col_hash{$self->{grid_solutions}->GetColLabelValue($index)} = $index;
	}

	my $budget = $self->{text_budget}->GetValue();
	
	
	#For each index being displayed calculate the rescaled tree
	foreach (keys %{$self->{index}}) {
		my $index = $self->{index}->{$_};
		if ( $index->{display} && !($index->{type} eq 'summary')) {
			
			#Delete any old rescaled trees 
			delete $index->{rescaled_tree};
			
			#Make a copy of the tree
			my $resized_tree = Bio::Phylo::IO->parse(
		   		'-string' => $self->{state}->{tree}->to_newick(),
   				'-format' => 'newick',
				)->first;	
	
			my %survival;
			if ($index->{type} eq 'solution') {
		
				#Get the species expenditure for this solution
				my %expenditure = $self->{state}->{solutions}->get_expenditure($index->{short});
				#Get the survival probabilities for this expenditure
				%survival = $self->{state}->{survival}->get_survival_probability(%expenditure);
				
			} elsif ($index->{type} eq 'species_specific') {
				my $lcIndex = lc($index->{short});
				my %index = %{eval('$self->{state}->{tree}->calc_'."$lcIndex")};
				%survival = $self->{state}->{survival}->get_survival_probability_budget_index($budget, %index);
				
			}
			
			#An all zero solution screws up plotting for some reason.
			#-- So don't plot all zero solutions
			my $allZero = 1;
			foreach (keys %survival) { if ($survival{$_}) {$allZero = 0;}}
			
			unless ($allZero) {	
				#Scale the tree
				$resized_tree->scale_survival(%survival);
				#Store it
				$index->{rescaled_tree} = $resized_tree;
				$index->{epd} = $resized_tree->calc_tree_length();
			} else {
				$index->{epd} = 0;
			}

			#Display epd and cost in grid	
			$self->{grid_solutions}->SetCellValue($nrows-1, $col_hash{$index->{short}}, $budget); 
			$self->{grid_solutions}->SetCellValue($nrows-2, $col_hash{$index->{short}}, $index->{epd}); 
		}
	}
}

#Sets the indices that can be added
sub set_combo_add_index {
	my $self = shift;
	#Clear combo box
	$self->{combo_add_index}->Clear();
	
	#Three arrays of indices to add, these will be separated by a blank entry
	my @indices_species_specific;
	my @indices_summary;
	my @indices_named;
	
	#Add indices to arrays
	foreach (keys %{$self->{index}}) {
		#Check that this index is not already displayed
		unless ($self->{index}->{$_}->{display} ) {
			#Check that it is not an index of the current budget
			unless ($self->{index}->{$_}->{type} eq 'solution' && $self->{index}->{$_}->{named} == 0) {
				eval 'push(@indices_'.$self->{index}->{$_}->{type}.', $self->{index}->{$_}->{long});';
			}
		}
	}
	
	#Sort the indices
	@indices_species_specific = sort @indices_species_specific;
	@indices_summary = sort @indices_summary;
	@indices_named = sort @indices_named;
	
	#Combine indices and add spacers
	my @combo_list = @indices_species_specific;
	if (@combo_list && (@indices_summary||@indices_named)) { push(@combo_list,''); }
	@combo_list = (@combo_list, @indices_summary);
	if (@indices_named && @combo_list && !($combo_list[-1] eq '') ) { push(@combo_list,''); }
	@combo_list = (@combo_list, @indices_named);
	
	#Add indices to the combo box
	foreach (@combo_list) {
		$self->{combo_add_index}->Append($_);
	}
	
	#If the combo box is empty disable the add button, otherwise enable.
	unless ($self->{combo_add_index}->GetCount() ) { 
		$self->{button_add_index}->Disable();
		$self->{combo_add_index}->Append('');
		$self->{combo_add_index}->SetValue('');
	}
	else { $self->{button_add_index}->Enable(); }
}

#Toggle the add index button appropriately
sub combo_selected {
	my $self = shift;
	my $index = $self->{combo_add_index}->GetStringSelection();
	if ($index) { $self->{button_add_index}->Enable(); }
	else {$self->{button_add_index}->Disable();}
}

#Add an index
sub add_index {
	my $self = shift;
	my $event = shift;

	my $index = $self->{combo_add_index}->GetStringSelection();
	
	#If no index was selected we update the combo box and return
	unless ($index) { 
		$self->set_combo_add_index();
		return;
	}
	
	#Otherwise select the index (if it is not already selected)
	unless ( $self->{index_r}->{$index}->{display} ) {
		$self->{index_r}->{$index}->{display} = 1;
		$self->set_combo_add_index();
		$self->update_grid();
	}

}

#Remove a column
sub remove_col {
	my $self = shift;

	my $select_index = $self->{grid_solutions}->GetSelectedCols;
	unless (defined $select_index) { return; }
	my $selection = $self->{grid_solutions}->GetColLabelValue($select_index);
	
	unless (defined $selection) { return; }

	if (defined $self->{index}->{$selection}) {
		$self->{index}->{$selection}->{display} = 0;
		$self->update_grid();
		$self->set_combo_add_index();
	}
	#$self->{grid_solutions}->EnableEditing();
	#	$self->{grid_solutions}->EnableCellEditControl();
	$self->{grid_solutions}->MoveCursorDown(0);
	$self->{grid_solutions}->MoveCursorRight(0);
}

#Set the tree
sub create_empty_grid {
	my $self = shift;
	
	unless (defined $self->{state}->{tree}) { return; }
	
	#Get an array of species names
	my @species;
	my @leaves = @{$self->{state}->{tree}->get_terminals};
	foreach (@leaves) { push(@species,$_->get_name()); }
	@species = sort @species;

	#Delete all rows & cols 
	$self->{grid_solutions}->DeleteRows(0, $self->{grid_solutions}->GetNumberRows);
	$self->{grid_solutions}->DeleteCols(0, $self->{grid_solutions}->GetNumberCols);

	#Add row labels
	if (defined $self->{grid_solutions}) {
		$self->{grid_solutions}->AppendRows(scalar @species+2);
		for (my $index = 0; $index < @species; $index++ ) { 
			$self->{grid_solutions}->SetRowLabelValue($index, $species[$index]);
		}
		$self->{grid_solutions}->SetRowLabelValue(scalar @species, 'Expected PD');
		$self->{grid_solutions}->SetRowLabelValue(scalar @species+1, 'Cost');		
	}
}

#we update the grid assuming that any columns in place contain correct details
sub update_grid {
	my $self = shift;

	unless (defined $self->{state}->{solutions}) { return; }

	#Get the indices that need to be added
	my %addition;
	my @deletion;

	#Get an array of species names
	my @species;
	my @leaves = @{$self->{state}->{tree}->get_terminals};
	foreach (@leaves) { push(@species,$_->get_name()); }
	@species = sort @species;
	
	#Get all indices that should be in the table
	foreach (keys %{$self->{index}}) {
		if ($self->{index}->{$_}->{display} ) { $addition{$_}=1; }
	}
	#Delete any keys that are already there
	for (my $colIndex = 0; $colIndex < $self->{grid_solutions}->GetNumberCols(); $colIndex++ ) {
		my $ind_key = $self->{grid_solutions}->GetColLabelValue($colIndex);
		if (defined $addition{$ind_key}) {
			delete $addition{$ind_key};
		} else {
			push(@deletion, $colIndex);
		}
	}

	#The mapping of species to row numbers
	my %species_mapping = $self->get_species_hash();
	
	#Add new columns for any of the indices not already there
	foreach (keys %addition) {
		#Get the index for the new column
		my $colIndex = $self->{grid_solutions}->GetNumberCols();
		#Append a new column and set its label
		$self->{grid_solutions}->AppendCols();
		$self->{grid_solutions}->SetColLabelValue($colIndex, $_);
		my %index;
		
		#Species specific indices
		if ($self->{index}->{$_}->{type} eq 'species_specific') {
			#Get the index values
			my $lcIndex = lc($_);
			%index = %{eval('$self->{state}->{tree}->calc_'."$lcIndex")};
			
		} elsif ($self->{index}->{$_}->{type} eq 'summary') {
			my $lcFunction = lc($_);
			%index = %{eval('$self->{state}->{solutions}->get_'."$lcFunction")};
			
			
		} elsif ($self->{index}->{$_}->{type} eq 'solution') {
			%index = $self->{state}->{solutions}->get_expenditure($_);
			 
		}
		
		#Set the index values
		foreach (@species) {
			my $value;
			if (defined $index{$_}) {
				$value = $index{$_};
			} else {
				$value = 0;				
			}
			if ($value == -1) {$value = 1; }
			print ':'.$value."\n";
			$self->{grid_solutions}->SetCellValue($species_mapping{$_}, $colIndex, sprintf("%.3g", $value));
		}
	
	}
	
	#Offset changes indices to account for previously deleted columns
	my $offset = 0;
	
	#Remove obsolete columns
	foreach (sort @deletion) {
		for (my $index = $_-$offset; $index < $self->{grid_solutions}->GetNumberCols()-1; $index++) {
			$self->{grid_solutions}->SetColLabelValue($index,  $self->{grid_solutions}->GetColLabelValue($index+1));
		}
		
		$self->{grid_solutions}->DeleteCols($_-$offset,1);
		$offset++;			
	}
	
	$self->{grid_solutions}->AutoSizeColumns();
	$self->calculate_scaled_trees();
}

#get a hash that maps species names to row numbers
sub get_species_hash {
	my $self = shift;
	my %output;
	for (my $rowIndex = 0; $rowIndex < $self->{grid_solutions}->GetNumberRows(); $rowIndex++ ) {
		$output{$self->{grid_solutions}->GetRowLabelValue($rowIndex)} = $rowIndex;
	}
	return %output;
}

#Displays a tree for the selected index
sub display_selected {
	my $self = shift;

	unless (defined $self->{state}->{solutions} && defined $self->{state}->{tree}) {
		#Should clear some stuff here....
		return;
	}
	
	#Get the actual solution object
	my $index = $self->get_selected_index();

	#Check if this tree has been plotted, if so we have nothing to do
	if (defined $self->{last_drawn_index} && defined $index->{rescaled_tree} && $self->{last_drawn_index} == $index->{rescaled_tree} ) { 
		return; 
	}
	
	$self->{last_drawn_index} = $index->{rescaled_tree};

	#Blank the treee
	$self->{wxtree}->set_tree();
	
	#No tree to draw	
	unless (defined $index->{rescaled_tree}) { return; }
	
	#Draw the tree
	$self->{wxtree}->set_tree($index->{rescaled_tree});
		
}

#Get the index currently selected
sub get_selected_index {
	my $self = shift;
	
	my $selection = $self->{grid_solutions}->GetSelectedCols();
	
	unless (defined $selection) { return; }

	#We store which of the indices this corresponds to
	my $index_label = $self->{grid_solutions}->GetColLabelValue($selection);

	return $self->{index}->{$index_label};
}


1;


